import numpy as np
from sklearn.cluster import KMeans
from src.common.logger import getLogger
from langchain_core.prompts import ChatPromptTemplate
from sklearn.metrics.pairwise import cosine_similarity
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class KMeanRAG:

    def __init__(self, llm_model, embed_model, vector_store, collection_prefix, library_number, n_clusters, top_k):
        self.llm_model = llm_model
        self.embed_model = embed_model
        self.vector_store = vector_store
        self.collection_prefix = collection_prefix
        self.library_number = library_number
        self.n_clusters = n_clusters
        self.top_k = top_k

    def calculate_document_similarity(self, query):
        logger.info("KMeanRAG calculate_document_similarity start")
        collection_name = self.collection_prefix + self.library_number
        doc_records, _ = self.vector_store.client.scroll(collection_name = collection_name, limit = 999999, offset = None)
        logger.info(f"KMeanRAG calculate_document_similarity doc_records len: {len(doc_records)}")
        documents = [doc.payload.get("page_content", "") for doc in doc_records]
        embed_query = self.embed_model.embed_query(query)
        vectors = self.embed_model.embed_documents(documents)
        similarity = cosine_similarity([embed_query], vectors)[0]
        logger.info(f"KMeanRAG calculate_document_similarity similarity len: {len(similarity)}")
        return documents, vectors, similarity

    def fit_kmean(self, vectors, similarity):
        logger.info("KMeanRAG fit_kmean start")
        top_n = min(self.top_k * 3, len(vectors))
        logger.info(f"KMeanRAG fit_kmean top_n: {top_n}")
        candidate_indices = np.argsort(similarity)[-top_n:][::-1]
        logger.info(f"KMeanRAG fit_kmean candidate_indices: {candidate_indices}")
        candidate_vectors = [vectors[indice] for indice in candidate_indices]

        n_clusters = min(self.n_clusters, len(candidate_vectors))
        logger.info(f"KMeanRAG fit_kmean n_clusters: {n_clusters}")
        kmeans = KMeans(n_clusters = n_clusters, random_state = 42, n_init = 1)
        clusters = kmeans.fit_predict(candidate_vectors)
        logger.info(f"KMeanRAG fit_kmean clusters: {clusters}")
        return clusters, candidate_indices

    def retrieve_kmean(self, documents, similarity, clusters, candidate_indices):
        logger.info("KMeanRAG retrieve_kmean start")
        cluster_dict = {}
        for idx, cluster_id in enumerate(clusters):
            cluster_nodes = cluster_dict.get(cluster_id, [])
            cluster_nodes.append((candidate_indices[idx]))
        logger.info(f"KMeanRAG retrieve_kmean cluster_dict len: {len(cluster_dict)}")

        select_indices = []
        for cluster_id in range(min(self.n_clusters, len(cluster_dict))):
            if cluster_dict.get(cluster_id):
                best_doc_indice = max(cluster_dict[cluster_id], key = lambda x: x[1])[0]
                select_indices.append(best_doc_indice)
        logger.info(f"KMeanRAG retrieve_kmean select_indices 11111: {select_indices}")

        remaining = [i for i in candidate_indices if i not in select_indices]
        logger.info(f"KMeanRAG retrieve_kmean remaining len: {len(remaining)}")
        if len(select_indices) < self.top_k and remaining:
            remaining_similarity =[similarity[i] for i in remaining]
            extra_indices = [remaining[i] for i in np.argsort(remaining_similarity)[(- self.top_k + len(select_indices)):]]
            select_indices.extend(extra_indices)
        logger.info(f"KMeanRAG retrieve_kmean select_indices 22222: {select_indices}")

        select_indices = select_indices[:self.top_k]
        logger.info(f"KMeanRAG retrieve_kmean select_indices 33333: {select_indices}")
        select_similarity = similarity[select_indices]
        logger.info(f"KMeanRAG retrieve_kmean select_similarity: {select_similarity}")
        return [documents[indice] for indice in select_indices]

    def invoke(self, query):
        logger.info("KMeanRAG invoke start")
        documents, vectors, similarity = self.calculate_document_similarity(query)
        clusters, candidate_indices = self.fit_kmean(vectors, similarity)
        retrieve_docs = self.retrieve_kmean(documents, similarity, clusters, candidate_indices)
        logger.info(f"KMeanRAG invoke retrieve_docs len: {len(retrieve_docs)}")

        template = """
            你是一位问题解答专家，请根据提供的上下文回答用户问题。
            
            上下文：{context}
            用户问题：{question}
            
            请必须用中文详尽的回答问题。
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        response = chain.invoke({ "context": retrieve_docs, "question": query })
        logger.info(f"KMeanRAG invoke response len: {len(response)}")
        return { "retrieve_docs": retrieve_docs, "chain_result": response }
